﻿// Copyright (c) 2012-2013 fancidev
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Drawing;
using System.ComponentModel;

namespace Util.Media
{
    public class OctreeColorQuantizer : ColorQuantizer
    {
        private OctreeNode root;
        private int NumColors; // number of distinct colors
        
        public OctreeColorQuantizer()
        {
        }

        public int DistinctColorCount { get { return NumColors; } }

        public override void Quantize(
            int[] pixels, int width, int height, int numColors)
        {
            BuildOctree(pixels, width, height, numColors);
        }

        public PaletteEntryInfo[] Palette { get; private set; }

        [Browsable(true)]
        public Color[] PaletteColors { get; private set; }

        [Browsable(true)]
        public int[] PaletteWeights { get; private set; }

        [Browsable(true)]
        public OctreeNode[] PaletteNodes { get; private set; }

        // TODO: The human eye is not sensitive to blue, so we can
        // reduce the precision in blue to allow for more precision
        // in the other color components.

        // Build an octree in the most straightforward way.
        private void BuildOctree(int[] argbArray, int width, int height, int target)
        {
            // Scan all pixels in the bitmap.
            UniformColorSampler sampler = new UniformColorSampler();
            sampler.Sample(argbArray, width, height);

            // Build an octree.
            root = OctreeNode.CreateRoot();
            List<OctreeNode> leaves = new List<OctreeNode>();
            foreach (KeyValuePair<int,int> kv in sampler.GetFrequencies())
            {
                Color c = Color.FromArgb(kv.Key);
                int frequency = kv.Value;

                // Traverse the tree from root to leaf one level by one level.
                OctreeNode node = root;
                for (int i = 7; i >= 0; i--)
                {
                    // Get a code between 0 and 7 inclusive for the i-th RGB layer.
                    int code = (((c.R >> i) & 1) << 2) |
                               (((c.G >> i) & 1) << 1) |
                               (((c.B >> i) & 1) << 0);

                    // Create the subtree if not present.
                    if (node.Children == null)
                        node.Children = new OctreeNode[8];
                    if (node.Children[code] == null)
                    {
                        node.Children[code] = new OctreeNode
                        {
                            Parent = node,
                            Color = Color.FromArgb(
                                c.R >> i << i,
                                c.G >> i << i,
                                c.B >> i << i),
                            Level = 8 - i,
                        };
                    }
                    node = node.Children[code];
                }

                // Now 'node' points to a leaf node.
                // Update the color count of this node.
                node.Mean = ColorCoordinates.FromColor(c);
                node.ColorCount = 1;
                node.PixelCount = (uint)frequency;
                leaves.Add(node);
                //node.R += c.R;
                //node.G += c.G;
                //node.B += c.B;
            }
            this.NumColors = leaves.Count;

            // Create a min-heap of octree nodes, where each node is ranked
            // by its color count.
#if false
            // Do a correction: let pixel-count' = log(pixel-count)/log(2)+1.
            // This can help to reduce the impact of large colors.
            // Note: doesn't work.
            foreach (OctreeNode node in leaves)
            {
                node.PixelCount = (uint)Math.Log(node.PixelCount, 2) + 1;
            }
#endif
            SortedSet<OctreeNode> heap = new SortedSet<OctreeNode>(
                leaves,
                new OctreeNodeComparer());

            // Merge the node with the least number of colors into its parent.
            // Repeat this process until there are no more than the requested
            // number of nodes left.
            while (heap.Count > target)
            {
                OctreeNode node = heap.Min;
                if (!heap.Remove(node))
                    throw new InvalidOperationException("Internal error: node not found.");

                // Remove the node's parent because we need to update its
                // count and color statistics.
                if (node.Parent == null)
                    continue;
                heap.Remove(node.Parent);

                // Also remove the parent's other children, because their
                // priority will change as well.
                foreach (OctreeNode child in node.Parent.Children)
                {
                    if (child != null && child.PixelCount > 0)
                        heap.Remove(child);
                }

                // Also remove this node's children, as their priority
                // will change as well.
                if (node.Children != null)
                {
                    foreach (OctreeNode child in node.Children)
                    {
                        if (child != null && child.PixelCount > 0)
                            heap.Remove(child);
                    }
                }

                // Update the parent's statistics, and then add back to the
                // priority queue.
                //node.Parent.R += node.R;
                //node.Parent.G += node.G;
                //node.Parent.B += node.B;
                node.Parent.Mean = ColorCoordinates.GetWeightedAverage(
                    node.Parent.Mean, node.Parent.PixelCount,
                    node.Mean, node.PixelCount);
                node.Parent.PixelCount += node.PixelCount;
                node.Parent.ColorCount += node.ColorCount;
                heap.Add(node.Parent);

                // Reset the attributes of the reduced node to prevent
                // double-counting if this node is added again later.
                node.PixelCount = 0;
                node.ColorCount = 0;
                //node.MeanR = 0;
                //node.MeanG = 0;
                //node.MeanB = 0;

                // Update the other children's statistics, and add them back
                // to the priority queue.
                foreach (OctreeNode child in node.Parent.Children)
                {
                    if (child != null && child.PixelCount > 0)
                        heap.Add(child);
                }

                // Also update this node's children, as their priority
                // will change as well.
                if (node.Children != null)
                {
                    foreach (OctreeNode child in node.Children)
                    {
                        if (child != null && child.PixelCount > 0)
                            heap.Add(child);
                    }
                }
            }

            // Get the representative colors of each node left.
            this.PaletteNodes = heap.Reverse().ToArray();
            this.PaletteColors = (from OctreeNode node in heap.Reverse()
                                  select node.AverageColor).ToArray();
            this.PaletteWeights = (from OctreeNode node in heap.Reverse()
                                   select (int)node.PixelCount).ToArray();

            // Quantize the input pixels.
            for (int i = 0; i < argbArray.Length; i++)
            {
#if true
                argbArray[i] = MapColor(Color.FromArgb(argbArray[i])).ToArgb();
#else
                int c = argbArray[i];
                bool found = false;

                // Reduce every three bits in (R,G,B) until we find an
                // entry in the palette.
                for (int j = 0; j < 8 && !found; j++)
                {
                    foreach (OctreeNode node in PaletteNodes)
                    {
                        if (node.Color.ToArgb() == c)
                        {
                            argbArray[i] = node.AverageColor.ToArgb();
                            found = true;
                            break;
                        }
                    }
                    c &= ~(0x00010101 << j);
                }

                // If we cannot find an opaque entry for this color, we have
                // to map it to the root color, provided that the root has not
                // been reduced.
                if (!found)
                {
                    //throw new InvalidOperationException(
                    //    "Cannot map color: " + argbArray[i].ToString("X8"));
                    // argbArray[i] = 0;
                    argbArray[i] = root.AverageColor.ToArgb();
                }
#endif
            }
        }

        public Color MapColor(Color c)
        {
            // Traverse the tree from the root element along the path
            // specified by the RGB layers in the given color. Return
            // the deepest node that has non-zero ColorCount, which
            // indicates that the node is in the palette.
            OctreeNode node = root;
            OctreeNode found = root;
            for (int i = 7; i >= 0; i--)
            {
                // Get a code between 0 and 7 inclusive for the i-th RGB layer.
                int code = (((c.R >> i) & 1) << 2) |
                           (((c.G >> i) & 1) << 1) |
                           (((c.B >> i) & 1) << 0);
                if (node.Children == null || node.Children[code] == null)
                    break;
                node = node.Children[code];
                if (node.ColorCount > 0)
                    found = node;
            }
            return found.AverageColor;
        }
    }

    [TypeConverter(typeof(ExpandableObjectConverter))]    
    public class OctreeNode //: PaletteEntryInfo
    {
        //private OctreeNode() { }

        /// <summary>
        /// Creates a root node that represents the entire color space.
        /// </summary>
        /// <returns></returns>
        public static OctreeNode CreateRoot()
        {
            return new OctreeNode();
        }

        public int Level { get; set; } // how deep is it? 0 is root, 8 is leaf
        public Color Color { get; set; } // representitative color of this node (cube)
        public ColorCoordinates Mean;
        //public uint R = 0; // sum of Red component of this node (cube).
        //public uint G = 0;
        //public uint B = 0;
        public uint ColorCount { get; set; } // nmber of distinct colors contained in this cube
        public uint PixelCount { get; set; } // number of pixels with this color
        public Color AverageColor
        {
            get { return ColorCoordinates.ToColor(Mean); }
            set { throw new NotImplementedException(); }
        }
        public int ChildCount
        {
            get
            {
                return Children == null ? 0 : Children.Length;
            }
        }
        public OctreeNode[] Children = null; // 8 children if this is not a leaf node
        public OctreeNode Parent = null; // parent node

        public override string ToString()
        {
            return string.Format(
                "L={0}, C={2}, P={1}",
                Level, PixelCount, ColorCount);
        }
    }

    class OctreeNodeComparer : IComparer<OctreeNode>
    {
        public int Compare(OctreeNode x, OctreeNode y)
        {
            if (x == null || y == null)
                throw new ArgumentNullException();

            // If one of the nodes is the root node, we always keep it
            // because otherwise we'll miss part of the color space.
            if (x.Level == 0 && y.Level == 0)
                return 0;
            if (x.Level == 0)
                return 1;
            if (y.Level == 0)
                return -1;
#if false
            if (x.Parent == null)
                return 1;
            if (y.Parent == null)
                return -1;
#endif
            int ret = 0;

#if true
            // Let's try some clever ranking... now we want to decide whether
            // to merge x into its parent or to merge y into its parent.
            // To evaluate this, we compute the extra SSE (sum square of error)
            // if we merge x with its parent.
            if (ret == 0)
            {
                float SSEx = GetExtraSSE(x);
                float SSEy = GetExtraSSE(y);
                ret = SSEx.CompareTo(SSEy);
#if false
                if (ret != 0)
                {
                    int kk = 1;
                }
#endif
            }

#else
            // Number of pixels weighted by number of colors.
            if (ret == 0)
                ret = (x.ColorCount * x.PixelCount).CompareTo
                      (y.ColorCount * y.PixelCount);

            // Then comes nodes with fewer distinct colors.
            if (ret == 0)
                ret = x.ColorCount.CompareTo(y.ColorCount);

            // Nodes with fewer pixels come first.
            if (ret == 0)
                ret = x.PixelCount.CompareTo(y.PixelCount);
#endif

            // For nodes with the same number of pixels, deeper nodes
            // come first.
            if (ret == 0)
                ret = y.Level.CompareTo(x.Level);

            // For two nodes with same number of pixels and at the same depth,
            // we must distinguish them to avoid SortedSet treating them as 
            // duplicates.
            if (ret == 0)
                ret = x.Color.ToArgb().CompareTo(y.Color.ToArgb());

            return ret;
        }

        private static float GetExtraSSE(OctreeNode x)
        {
            OctreeNode y = x.Parent;
            if (x.PixelCount == 0 || y.PixelCount == 0)
                return 0;

            float n1 = x.PixelCount;
            float n2 = y.PixelCount;
            return ColorCoordinates.DistanceSquared(x.Mean, y.Mean) * n1 * n2 / (n1 + n2);
        }
    }
}
